#include "sandbox_manager.h"
#include <sys/wait.h>
#include <sys/resource.h>
#include <unistd.h>
#include <cstring>
#include <errno.h>

namespace judge {

SandboxManager::SandboxManager() {
    // 初始化日志
    spdlog::set_level(spdlog::level::debug);

    // 初始化其他组件
    namespace_manager_ = std::make_unique<NamespaceManager>();
    cgroup_manager_ = std::make_unique<CgroupManager>();
}

int SandboxManager::run() {
    try {
        spdlog::debug("Starting sandbox execution");

        // 初始化命名空间
        namespace_manager_->initialize();

        // 创建子进程
        return executeCommand(getpid());
    } catch (const std::exception& e) {
        spdlog::error("Sandbox execution failed: {}", e.what());
        return -1;
    }
}

int SandboxManager::executeCommand(pid_t) {
    spdlog::debug("Executing command: {}", command_);

    // 创建管道用于重定向输出
    int stdout_pipe[2], stderr_pipe[2];
    if (pipe(stdout_pipe) == -1 || pipe(stderr_pipe) == -1) {
        throw std::runtime_error("Failed to create pipes");
    }

    // 创建子进程
    pid_t child_pid = fork();
    if (child_pid == -1) {
        throw std::runtime_error("Fork failed");
    }

    if (child_pid == 0) {
        // 子进程
        try {
            // 重定向标准输出和标准错误
            dup2(stdout_pipe[1], STDOUT_FILENO);
            dup2(stderr_pipe[1], STDERR_FILENO);
            close(stdout_pipe[0]);
            close(stdout_pipe[1]);
            close(stderr_pipe[0]);
            close(stderr_pipe[1]);

            // 设置工作目录
            if (!working_dir_.empty() && chdir(working_dir_.c_str()) == -1) {
                throw std::runtime_error("Failed to change working directory");
            }

            // 设置资源限制
            setResourceLimits();

            // 执行命令
            std::vector<char*> args;
            std::istringstream iss(command_);
            std::string token;
            while (iss >> token) {
                char* arg = new char[token.length() + 1];
                strcpy(arg, token.c_str());
                args.push_back(arg);
            }
            args.push_back(nullptr);

            spdlog::debug("Executing command with {} arguments", args.size() - 1);
            execvp(args[0], args.data());

            // 如果execvp失败
            throw std::runtime_error("Failed to execute command: " + std::string(strerror(errno)));
        } catch (const std::exception& e) {
            spdlog::error("Child process failed: {}", e.what());
            _exit(1);
        }
    }

    // 父进程
    close(stdout_pipe[1]);
    close(stderr_pipe[1]);

    // 读取输出
    std::thread stdout_thread([this, &stdout_pipe]() {
        char buffer[4096];
        ssize_t count;
        while ((count = read(stdout_pipe[0], buffer, sizeof(buffer))) > 0) {
            stdout_.append(buffer, count);
        }
    });

    std::thread stderr_thread([this, &stderr_pipe]() {
        char buffer[4096];
        ssize_t count;
        while ((count = read(stderr_pipe[0], buffer, sizeof(buffer))) > 0) {
            stderr_.append(buffer, count);
        }
    });

    // 等待输出读取完成
    stdout_thread.join();
    stderr_thread.join();

    // 关闭管道
    close(stdout_pipe[0]);
    close(stderr_pipe[0]);

    // 等待子进程结束
    int status;
    waitpid(child_pid, &status, 0);

    // 检查退出状态
    if (WIFEXITED(status)) {
        int exit_code = WEXITSTATUS(status);
        spdlog::debug("Process exited normally with code {}", exit_code);
        return exit_code;
    } else if (WIFSIGNALED(status)) {
        int signal = WTERMSIG(status);
        spdlog::debug("Process terminated by signal {}", signal);
        return 128 + signal;
    }

    spdlog::debug("Process terminated abnormally");
    return -1;
}

void SandboxManager::setResourceLimits() {
    spdlog::debug("Setting resource limits");

    // 设置CPU时间限制
    if (cpu_time_limit_ > 0) {
        struct rlimit rlim;
        rlim.rlim_cur = rlim.rlim_max = cpu_time_limit_;
        if (setrlimit(RLIMIT_CPU, &rlim) == -1) {
            throw std::runtime_error("Failed to set CPU time limit");
        }
    }

    // 设置内存限制
    if (memory_limit_ > 0) {
        struct rlimit rlim;
        rlim.rlim_cur = rlim.rlim_max = memory_limit_ * 1024 * 1024; // 转换为字节
        if (setrlimit(RLIMIT_AS, &rlim) == -1) {
            throw std::runtime_error("Failed to set memory limit");
        }
    }

    // 设置输出大小限制
    if (output_limit_ > 0) {
        struct rlimit rlim;
        rlim.rlim_cur = rlim.rlim_max = output_limit_ * 1024; // 转换为字节
        if (setrlimit(RLIMIT_FSIZE, &rlim) == -1) {
            throw std::runtime_error("Failed to set output size limit");
        }
    }

    spdlog::debug("Resource limits set successfully");
} 

} // namespace judge 